/*==============================================================================
案例1：员工流失预测 - 模型训练与评估
================================================================================

业务目标：
1. 预测哪些员工可能离职
2. 识别影响员工离职的关键因素
3. 为人力资源部门提供决策支持

模型方法：
- Logistic 回归（基准模型）
- Lasso Logit（特征选择）
- 随机森林（非线性模型）

评估指标：
- 准确率、精确率、召回率
- ROC 曲线和 AUC
- 混淆矩阵

作者：Stata ML Course
日期：2025-11-03
==============================================================================*/

clear all
set more off
capture log close
log using "output/cases/case01_modeling.log", replace

*------------------------------------------------------------------------------
* 1. 加载数据
*------------------------------------------------------------------------------

display "=" * 80
display "案例1：员工流失预测 - 模型训练"
display "=" * 80
display ""

use "data/cases/employee_attrition.dta", clear

display "数据概览："
describe
display ""
display "样本量: " _N

*------------------------------------------------------------------------------
* 2. 数据分割
*------------------------------------------------------------------------------

display ""
display "数据分割（80% 训练集，20% 测试集）"
display "-" * 80

set seed 20251103
gen random = runiform()
gen train = (random <= 0.8)
gen test = (random > 0.8)

tab train
display "训练集样本: " _N if train
display "测试集样本: " _N if test

*------------------------------------------------------------------------------
* 3. 模型1：Logistic 回归（基准模型）
*------------------------------------------------------------------------------

display ""
display "模型1：Logistic 回归"
display "=" * 80

* 训练模型
logit attrition wage hours ttl_exp tenure age grade married union south ///
    if train

* 保存模型
estimates store logit_model

* 显示结果
display ""
display "模型系数："
logit

* 边际效应
display ""
display "平均边际效应："
margins, dydx(*) post
estimates store logit_margins

* 预测
estimates restore logit_model
predict prob_logit, pr
gen pred_logit = (prob_logit > 0.5)

* 评估（测试集）
display ""
display "测试集评估："
display "-" * 80

* 混淆矩阵
tab attrition pred_logit if test, row col

* 准确率
gen correct_logit = (attrition == pred_logit) if test
summarize correct_logit
local acc_logit = r(mean)
display "准确率: " %6.4f `acc_logit'

* 计算精确率和召回率
count if pred_logit == 1 & attrition == 1 & test
local tp = r(N)
count if pred_logit == 1 & attrition == 0 & test
local fp = r(N)
count if pred_logit == 0 & attrition == 1 & test
local fn = r(N)

local precision = `tp' / (`tp' + `fp')
local recall = `tp' / (`tp' + `fn')
local f1 = 2 * `precision' * `recall' / (`precision' + `recall')

display "精确率: " %6.4f `precision'
display "召回率: " %6.4f `recall'
display "F1分数: " %6.4f `f1'

* ROC 曲线
estimates restore logit_model
lroc if test, title("Logistic 回归 ROC 曲线")
graph export "output/cases/figures/case01_logit_roc.png", replace

*------------------------------------------------------------------------------
* 4. 模型2：Lasso Logit（特征选择）
*------------------------------------------------------------------------------

display ""
display "模型2：Lasso Logistic 回归"
display "=" * 80

* 训练 Lasso Logit
lasso logit attrition wage hours ttl_exp tenure age grade married union south ///
    monthly_income exp_tenure_ratio is_new_employee is_veteran is_overwork ///
    if train, selection(cv) rseed(20251103)

* 查看选择的特征
display ""
display "Lasso 选择的特征："
lassocoef

* 预测
predict prob_lasso, pr
gen pred_lasso = (prob_lasso > 0.5)

* 评估（测试集）
display ""
display "测试集评估："
display "-" * 80

tab attrition pred_lasso if test, row col

gen correct_lasso = (attrition == pred_lasso) if test
summarize correct_lasso
local acc_lasso = r(mean)
display "准确率: " %6.4f `acc_lasso'

*------------------------------------------------------------------------------
* 5. 模型3：随机森林（如果已安装）
*------------------------------------------------------------------------------

display ""
display "模型3：随机森林"
display "=" * 80

* 检查是否安装 rforest
capture which rforest
if _rc == 0 {
    display "训练随机森林模型..."
    
    * 训练随机森林
    rforest attrition wage hours ttl_exp tenure age grade married union south ///
        if train, type(class) iter(100) seed(20251103)
    
    display "随机森林训练完成"
    display "注意：rforest 的预测功能有限，建议使用 Logit 或 Lasso"
}
else {
    display "未安装 rforest 包"
    display "可以运行: ssc install rforest"
}

*------------------------------------------------------------------------------
* 6. 模型对比
*------------------------------------------------------------------------------

display ""
display "模型对比"
display "=" * 80

* 创建对比表
matrix results = J(2, 3, .)
matrix rownames results = "Logit" "Lasso"
matrix colnames results = "Accuracy" "Precision" "Recall"

* Logit 结果
matrix results[1, 1] = `acc_logit'
matrix results[1, 2] = `precision'
matrix results[1, 3] = `recall'

* Lasso 结果（重新计算）
count if pred_lasso == 1 & attrition == 1 & test
local tp_lasso = r(N)
count if pred_lasso == 1 & attrition == 0 & test
local fp_lasso = r(N)
count if pred_lasso == 0 & attrition == 1 & test
local fn_lasso = r(N)

local precision_lasso = `tp_lasso' / (`tp_lasso' + `fp_lasso')
local recall_lasso = `tp_lasso' / (`tp_lasso' + `fn_lasso')

matrix results[2, 1] = `acc_lasso'
matrix results[2, 2] = `precision_lasso'
matrix results[2, 3] = `recall_lasso'

display ""
display "模型性能对比："
matrix list results

*------------------------------------------------------------------------------
* 7. 特征重要性分析
*------------------------------------------------------------------------------

display ""
display "特征重要性分析"
display "=" * 80

* 使用 Logit 模型的边际效应
estimates restore logit_margins

* 绘制系数图
coefplot, drop(_cons) xline(0) ///
    title("特征重要性（边际效应）") ///
    xtitle("对离职概率的影响") ///
    scheme(s2color)
graph export "output/cases/figures/case01_feature_importance.png", replace

*------------------------------------------------------------------------------
* 8. 业务洞察可视化
*------------------------------------------------------------------------------

display ""
display "生成业务洞察图表..."
display "-" * 80

* 8.1 预测概率分布
twoway (histogram prob_logit if attrition == 0 & test, color(blue%30)) ///
       (histogram prob_logit if attrition == 1 & test, color(red%30)), ///
       legend(label(1 "未离职") label(2 "已离职")) ///
       title("预测离职概率分布") ///
       xtitle("预测概率") ytitle("频数") ///
       scheme(s2color)
graph export "output/cases/figures/case01_prob_distribution.png", replace

* 8.2 高风险员工识别
gen risk_level = 1 if prob_logit < 0.3
replace risk_level = 2 if prob_logit >= 0.3 & prob_logit < 0.7
replace risk_level = 3 if prob_logit >= 0.7 & prob_logit < .
label define risk_lbl 1 "低风险" 2 "中风险" 3 "高风险"
label values risk_level risk_lbl

graph bar (count), over(risk_level) ///
    title("员工离职风险分布") ///
    ytitle("人数") ///
    blabel(bar, format(%9.0f)) ///
    scheme(s2color)
graph export "output/cases/figures/case01_risk_distribution.png", replace

* 8.3 不同工资水平的离职风险
egen wage_quartile = cut(wage), group(4)
graph bar (mean) prob_logit, over(wage_quartile) ///
    title("不同工资水平的平均离职风险") ///
    ytitle("平均离职概率") ///
    blabel(bar, format(%9.3f)) ///
    scheme(s2color)
graph export "output/cases/figures/case01_wage_risk.png", replace

*------------------------------------------------------------------------------
* 9. 生成业务报告
*------------------------------------------------------------------------------

display ""
display "生成业务报告..."
display "-" * 80

* 识别高风险员工
preserve
keep if test & risk_level == 3
sort prob_logit
list wage tenure age grade prob_logit in 1/10, clean

display ""
display "高风险员工特征："
summarize wage tenure age grade
restore

*------------------------------------------------------------------------------
* 10. 保存结果
*------------------------------------------------------------------------------

* 保存预测结果
preserve
keep if test
keep attrition prob_logit pred_logit risk_level wage tenure age grade
export delimited "output/cases/case01_predictions.csv", replace
restore

* 保存完整数据
save "data/cases/employee_attrition_results.dta", replace

display ""
display "=" * 80
display "建模完成！"
display "=" * 80
display ""
display "模型性能："
display "  Logit 准确率: " %6.2f `acc_logit' * 100 "%"
display "  Lasso 准确率: " %6.2f `acc_lasso' * 100 "%"
display ""
display "输出文件："
display "  - 预测结果: output/cases/case01_predictions.csv"
display "  - 图表: output/cases/figures/case01_*.png"
display "  - 日志: output/cases/case01_modeling.log"

log close

